#!/usr/bin/env python
# -*- coding: UTF-8 -*-
import os
import time
import requests
import argparse
import sys
import logging
from pathlib import Path
from typing import Dict, List, Optional, Union, Any


# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger('apk_reinforce')


class DownloadHandler:
    def __init__(self, url: str, username: str, password: str):
        self.url = url
        self.username = username
        self.password = password
        self.download_path = "./downloads"
        
    def setup_download_dir(self):
        """创建下载目录"""
        if not os.path.exists(self.download_path):
            os.makedirs(self.download_path)
            logger.info(f"创建下载目录: {self.download_path}")
            
    def download_apk(self) -> Optional[str]:
        """下载APK文件"""
        try:
            self.setup_download_dir()
            response = requests.get(
                self.url,
                auth=(self.username, self.password),
                stream=True,
                timeout=30
            )
            response.raise_for_status()
            
            # 从URL中提取文件名，如果没有则使用默认名
            file_name = self.url.split('/')[-1] or "app.apk"
            file_path = os.path.join(self.download_path, file_name)
            
            with open(file_path, 'wb') as f:
                for chunk in response.iter_content(chunk_size=8192):
                    if chunk:
                        f.write(chunk)
                        
            logger.info(f"APK下载完成: {file_path}")
            return file_path
            
        except Exception as e:
            logger.error(f"下载APK失败: {str(e)}")
            sys.exit(1)


class LocalApkHandler:
    def __init__(self, local_path: str):
        self.local_path = local_path
        self.download_path = "./downloads"
        
    def setup_download_dir(self):
        """创建下载目录"""
        if not os.path.exists(self.download_path):
            os.makedirs(self.download_path)
            logger.info(f"创建下载目录: {self.download_path}")
            
    def process_apk(self) -> Optional[str]:
        """处理本地APK文件"""
        try:
            # 验证本地文件是否存在
            if not os.path.exists(self.local_path):
                logger.error(f"本地APK文件不存在: {self.local_path}")
                sys.exit(1)
                
            # 如果文件不在下载目录中，复制到下载目录
            self.setup_download_dir()
            file_name = os.path.basename(self.local_path)
            file_path = os.path.join(self.download_path, file_name)
            
            if self.local_path != file_path:
                import shutil
                shutil.copy2(self.local_path, file_path)
                logger.info(f"本地APK文件已复制到: {file_path}")
            else:
                logger.info(f"使用本地APK文件: {file_path}")
                
            return file_path
            
        except Exception as e:
            logger.error(f"处理本地APK文件失败: {str(e)}")
            sys.exit(1)


class ReinforceHandler:
    def __init__(self, base_url: str = "http://isec.iflytek.com:8000/"):
        self.base_url = base_url
        self.syncore = "E54BB07E4ECE46AAB63B1C9D3C3203B9"
        self.expire = 30
        self.access_token: Optional[str] = None
        
    def _process_response(self, response: requests.Response, error_msg: str) -> Dict[str, Any]:
        """处理API响应，统一错误处理"""
        try:
            # 检查HTTP错误
            response.raise_for_status()
            
            # 解析JSON响应
            resp_json = response.json()
            
            # 检查业务逻辑错误
            if resp_json.get('code') == '508':
                logger.error(f"加固服务异常: {resp_json}")
                sys.exit(1)
                
            return resp_json
        except requests.exceptions.HTTPError as e:
            logger.error(f"{error_msg}: HTTP错误 - {str(e)}")
            sys.exit(1)
        except ValueError as e:
            logger.error(f"{error_msg}: JSON解析错误 - {str(e)}")
            logger.error(f"原始响应内容: {response.text}")
            sys.exit(1)
        except Exception as e:
            logger.error(f"{error_msg}: {str(e)}")
            sys.exit(1)
        
    def get_access_token(self) -> Optional[str]:
        """获取访问令牌"""
        try:
            headers = {'Content-Type': 'application/json ;charset=UTF-8'}
            params = {'clientKey': self.syncore, 'expire': self.expire}
            response = requests.post(
                f"{self.base_url}sec/open/auth/accessToken",
                headers=headers,
                json=params
            )
            
            # 处理响应
            resp_json = self._process_response(response, "获取访问令牌失败")
            
            # 检查返回结果是否包含obj字段和accessToken
            if not resp_json.get('obj') or 'accessToken' not in resp_json.get('obj', {}):
                logger.error(f"返回数据格式异常: {resp_json}")
                sys.exit(1)
                
            self.access_token = resp_json['obj']['accessToken']
            return self.access_token
        except Exception as e:
            logger.error(f"获取访问令牌失败: {str(e)}")
            sys.exit(1)
            
    def upload_apk(self, file_path: str) -> Optional[Dict[str, Any]]:
        """上传APK文件"""
        try:
            headers = {"Authorization": f"Bearer {self.access_token}"}
            files = {
                'file': (
                    os.path.basename(file_path),
                    open(file_path, 'rb'),
                    'application/vnd.android.package-archive'
                )
            }
            response = requests.post(
                f"{self.base_url}sec/open/file/upload",
                headers=headers,
                files=files
            )
            
            # 处理响应
            resp_json = self._process_response(response, "上传APK失败")
            
            # 检查返回结果是否包含obj字段
            if not resp_json.get('obj'):
                logger.error(f"返回数据格式异常: {resp_json}")
                sys.exit(1)
            
            logger.info(f"APK文件上传成功: {os.path.basename(file_path)}")
            return resp_json['obj']
        except Exception as e:
            logger.error(f"上传APK失败: {str(e)}")
            sys.exit(1)
            
    def create_reinforce_task(self, upload_result: Dict[str, Any]) -> Optional[int]:
        """创建加固任务"""
        try:
            headers = {
                'Content-Type': 'application/json ;charset=UTF-8',
                "Authorization": f"Bearer {self.access_token}"
            }
            params = {
                "dataSrc": "SYNCORE",
                "fileName": upload_result['fileName'],
                "downloadName": upload_result['downloadName'],
                "fileSaveName": upload_result['fileSaveName'],
                "filePath": upload_result['filePath'],
                "minioFileUrl": upload_result['minioFileUrl'],
                "md5": upload_result['md5'],
                "userAccount": "xbli16",
                "keystoreId": 35
            }
            response = requests.post(
                f"{self.base_url}sec/open/reinforceCommon/addTask",
                headers=headers,
                json=params
            )
            
            # 处理响应
            resp_json = self._process_response(response, "创建加固任务失败")
            
            # 检查返回结果是否包含obj字段和taskId
            if not resp_json.get('obj') or 'taskId' not in resp_json.get('obj', {}):
                logger.error(f"返回数据格式异常: {resp_json}")
                sys.exit(1)
                
            task_id = resp_json['obj']['taskId']
            logger.info(f"加固任务创建成功，任务ID: {task_id}")
            return task_id
        except Exception as e:
            logger.error(f"创建加固任务失败: {str(e)}")
            sys.exit(1)
            
    def check_task_status(self, task_id: int) -> Optional[Dict[str, Any]]:
        """查询任务状态"""
        try:
            headers = {
                'Content-Type': 'application/json ;charset=UTF-8',
                "Authorization": f"Bearer {self.access_token}"
            }
            response = requests.post(
                f"{self.base_url}sec/open/reinforceCommon/taskStatus",
                headers=headers,
                json=[task_id]
            )
            
            # 处理响应
            resp_json = self._process_response(response, "查询任务状态失败")
            
            # 检查返回结果是否包含obj字段且不为空
            if not resp_json.get('obj') or not isinstance(resp_json['obj'], list) or len(resp_json['obj']) == 0:
                logger.error(f"返回数据格式异常: {resp_json}")
                sys.exit(1)
                
            # 返回任务状态
            return resp_json["obj"][0]
        except Exception as e:
            logger.error(f"查询任务状态失败: {str(e)}")
            sys.exit(1)
            
    def download_reinforced_apk(self, task_id: int) -> Optional[str]:
        """下载加固后的APK"""
        try:
            headers = {"Authorization": f"Bearer {self.access_token}"}
            response = requests.get(
                f"{self.base_url}sec/open/reinforceCommon/download?taskId={task_id}",
                headers=headers,
                stream=True,
                timeout=30
            )
            
            # 检查HTTP错误
            response.raise_for_status()
            
            # 检查响应头的Content-Type，确保是APK文件
            content_type = response.headers.get('Content-Type', '')
            if 'application/json' in content_type:
                # 如果是JSON响应，可能是错误信息
                try:
                    resp_json = response.json()
                    if resp_json.get('code') == '508':
                        logger.error(f"加固服务异常: {resp_json}")
                        sys.exit(1)
                    logger.error(f"下载失败，服务器返回: {resp_json}")
                    sys.exit(1)
                except Exception:
                    logger.error("下载失败，无法解析响应内容")
                    sys.exit(1)
            
            file_name = f"reinforced_{task_id}.apk"
            file_path = os.path.join("./downloads", file_name)
            
            with open(file_path, 'wb') as f:
                for chunk in response.iter_content(chunk_size=8192):
                    if chunk:
                        f.write(chunk)
                        
            logger.info(f"加固后的APK下载完成: {file_path}")
            return file_path
        except Exception as e:
            logger.error(f"下载加固后的APK失败: {str(e)}")
            sys.exit(1)


def parse_args() -> argparse.Namespace:
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description='APK安全加固工具')
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument('--remote-apk', help='APK在制品库中的路径，例如：appWarehouse-AIONA19XM_SOP7/sop7/gxatek/SYBtService/SYBtService.apk')
    group.add_argument('--local-apk', help='本地APK文件的路径')
    parser.add_argument('--reinforce', type=str, choices=['true', 'false'], default='true',
                      help='是否进行加固，可选值：true/false，默认为true')
    return parser.parse_args()


def main():
    # 解析命令行参数
    args = parse_args()
    
    # 初始化加固处理器
    reinforcer = ReinforceHandler()
    
    try:
        # 处理APK文件获取
        apk_path: Optional[str] = None

        if args.reinforce.lower() == 'false':
            logger.info("当前版本已有加固，跳过加固步骤")
            sys.exit(0)
        
        if args.remote_apk:
            # 构建完整的下载URL
            base_url = "http://console.devops.gxatek.com/restapi/ipackage/v1/download/artifact?fullPath=/XHZL-private-repo/"
            full_url = f"{base_url}{args.remote_apk}"
            logger.info(f"准备从远程下载APK: {args.remote_apk}")
            
            # 配置下载参数
            download_config = {
                "url": full_url,
                "username": "ysliang",
                "password": "AKCp8mYeP29WkmkSvPGcyBP3mDZyk9kJFgx1d2jmAmL6phkBZYqVB13G9jFbUh3K8G1KKY219"
            }
            
            # 下载APK
            downloader = DownloadHandler(**download_config)
            apk_path = downloader.download_apk()
        else:
            # 处理本地APK
            logger.info(f"准备处理本地APK: {args.local_apk}")
            local_handler = LocalApkHandler(args.local_apk)
            apk_path = local_handler.process_apk()
            
        if not apk_path:
            logger.error("无法获取APK文件路径")
            sys.exit(1)
        
        # 2. 获取访问令牌
        logger.info("正在获取访问令牌...")
        if not reinforcer.get_access_token():
            logger.error("获取访问令牌失败")
            sys.exit(1)
            
        # 3. 上传APK并创建加固任务
        logger.info("正在上传APK文件...")
        upload_result = reinforcer.upload_apk(apk_path)
        if not upload_result:
            logger.error("上传APK文件失败")
            sys.exit(1)
            
        logger.info("正在创建加固任务...")
        task_id = reinforcer.create_reinforce_task(upload_result)
        if not task_id:
            logger.error("创建加固任务失败")
            sys.exit(1)
            
        # 4. 轮询任务状态并下载加固后的APK
        logger.info(f"开始轮询加固任务状态，任务ID: {task_id}")
        max_attempts = 30  # 最大尝试次数，防止无限循环
        attempt = 0
        while attempt < max_attempts:
            task_status = reinforcer.check_task_status(task_id)
            if not task_status:
                logger.error("查询任务状态失败")
                sys.exit(1)
                
            if task_status["downloadApk"] and task_status["taskComplete"]:
                logger.info("加固任务完成，准备下载加固后的APK")
                reinforcer.download_reinforced_apk(task_id)
                break
                
            logger.info(f"任务状态: {task_status}")
            time.sleep(20)
            attempt += 1
            
        if attempt >= max_attempts:
            logger.error("任务超时，未能完成加固")
            sys.exit(1)
            
    except Exception as e:
        logger.error(f"程序执行出错: {str(e)}")
        sys.exit(1)


if __name__ == '__main__':
    main()
